import h5py
import numpy as np
import pandas as pd
import wntr
import matplotlib.pyplot as plt
wn = wntr.network.WaterNetworkModel(r".\覆盖空间\基础.inp")

def perception(C,I):
    i = I[C,:][:,C]
    rate = i.sum(axis=0)/i.shape[0]
    return rate

def imatrix(delta=None,time=None):
    # 压力阈值
    df = pd.read_excel(r"pressure.xlsx",header=0)
    
    Is = []
    if delta == None:
        delta = [0.25, 0.5, 0.75, 1]
    for d in delta:
        file = f"苏州60天均时_P_{d}.mat"
        dt = h5py.File(file)
        data = np.array(dt['P_abnormal'])
        data = data.T

        ## 不同时间
        if time == None:
            time = [2,4,6,8,10,12,14,16,18,20,22]
        for h in time:
            bp = data[h::25, :]
            I = bp - df['99%'].values[h * 60]+1.6
            I[I > 0] = 0
            I[I < 0] = 1
            Is.append(I)
    return Is

def layout():

    # 指示矩阵
    ## 不同程度
    Is = imatrix()

    area = pd.read_excel(r"E:\用户\桌面\研究内容\大论文\代码\model\cluster\run\area.xlsx",header=0,sheet_name='area')

    nc=[]
    for n in range(1,21):
        nc.append(list(area[area['SDCN_class']==n]['node_index'].values))

    ms=[]
    r=[]
    atb = np.zeros(wn.num_nodes)
    atb[area['node_index'].values] = area['SDCN_class'].values
    atr=dict(zip(wn.node_name_list,atb))
    for cls in nc:
        print('nodes:%d'%len(cls))
        rate=0
        for j in Is:
            rate += perception(cls,j)
        rate = rate/len(Is)
        m=[]
        for k in np.where(rate==rate.max())[0]:
            m.append(cls[k])
            atr[wn.node_name_list[cls[k]]]=50
        r.append(max(rate))
        print('monitors:%d'%len(m))
        ms.append(m)
    print(ms)
    print(r)
    popup = pd.DataFrame()
    popup['node'] = wn.node_name_list
    popup['index'] = np.arange(467)
    pops = popup.set_index('node')

    wntr.graphics.plot_interactive_network(wn,atr,add_to_node_popup=pops)


def test_monitor(m,delta,time):
    p=[]
    Is = imatrix(delta,time)
    print(len(Is))
    for n,i in enumerate(Is,start=1):
        print('工况%d：'%n)
        nodes = []
        for j in m:
            nodes = nodes + list(np.where(i[:,j]==1)[0])
        print('覆盖率：%.6f'%(len(set(nodes))/i.shape[0]))
        p.append(len(set(nodes))/i.shape[0])
        # print('未覆盖的节点：%s'%(set(np.arange(0,464))-set(nodes)))
        # wn = wntr.network.WaterNetworkModel(r"E:\用户\桌面\研究内容\我的小论文\思路\生成数据\爆管压力\聚类\覆盖空间\基础.inp")
        # atr=dict(zip(wn.node_name_list,np.zeros(wn.num_nodes)))
        # for t in set(np.arange(0,464))-set(nodes):
            # atr[wn.node_name_list[t]] = 1
            # node = wn.get_node(wn.node_name_list[t])
            # node.elevation = 10
        # wn.write_inpfile(f'E:\用户\桌面\研究内容\我的小论文\思路\生成数据\爆管压力\聚类\覆盖空间\sp{len(m)}_{n}.inp')
        # wntr.graphics.plot_network(wn,atr,add_colorbar=False)
        # plt.savefig(f'E:\用户\桌面\研究内容\我的小论文\思路\生成数据\爆管压力\聚类\覆盖空间\sp{len(m)}_{n}.png')
            
    return p


if __name__ == '__main__':
    sheet = ['monitor_6','monitor_8','monitor_10','monitor_12','monitor_20']
    for s in sheet:
        df = pd.read_excel(r"E:\用户\桌面\研究内容\大论文\代码\model\cluster\monitor-scheme.xlsx",sheet_name=s)
        atb = np.zeros(wn.num_nodes)
        atb[df['node_index'].values] = df['SDCN_class'].values
        atr=dict(zip(wn.node_name_list,atb))
        popup = pd.DataFrame()
        popup['node'] = wn.node_name_list
        popup['index'] = np.arange(467)
        pops = popup.set_index('node')
        wntr.graphics.plot_interactive_network(wn,atr,node_attribute_name='分区',add_to_node_popup=pops,figsize=[600,740],filename=s+'.html')